﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Windows.Media;
using System.Windows.Media.Imaging;

namespace ArgusMedia.Video
{
	public enum Y4MColorSpace
	{
		//C420jpeg,
		//C420paldv,
		C420,
		C422,
		C444
	}

	public struct Rational
	{
		public int Numinator { get; set; }
		public int Denominator { get; set; }

		public Rational(int Num, int Den)
			:this()
		{
			this.Numinator = Num;
			this.Denominator = Den;
		}
	}

	/// <summary>
	/// http://wiki.multimedia.cx/index.php?title=YUV4MPEG2
	/// </summary>
	public class Y4MEncoder
	{
		#region Static Helpers
		static readonly byte[] HeaderFile = ToByteArray("YUV4MPEG2");
		static readonly byte[] HeaderFrame = ToByteArray("FRAME");
		static readonly byte[] ParaWidth = ToByteArray(" W");
		static readonly byte[] ParaHeight = ToByteArray(" H");
		static readonly byte[] ParaFramerate = ToByteArray(" F");

		static byte[] ToByteArray(string s)
		{
			byte[] data = new byte[s.Length];
			for (int i = 0; i < s.Length; i++)
				data[i] = (byte)s[i];
			return data;
		}

		private static void WriteHeader(Stream stream, int Width, int Height, Rational Framerate, Y4MColorSpace Colorspace)
		{
			WriteFileheader(stream);
			WriteParameterWidth(stream, Width);
			WriteParameterHeight(stream, Height);
			WriteParameterFramerate(stream, Framerate);
			WriteParameterColorspace(stream, Colorspace);
		}

		private static void WriteFileheader(Stream stream)
		{
			stream.Write(HeaderFile, 0, HeaderFile.Length);
		}

		private static void WriteParameterWidth(Stream stream, int Width)
		{
			List<byte> data = new List<byte>();
			data.AddRange(ParaWidth);
			data.AddRange(ToByteArray(Width.ToString()));
			stream.Write(data.ToArray(), 0, data.Count);
		}

		private static void WriteParameterHeight(Stream stream, int Height)
		{
			List<byte> data = new List<byte>();
			data.AddRange(ParaHeight);
			data.AddRange(ToByteArray(Height.ToString()));
			stream.Write(data.ToArray(), 0, data.Count);
		}

		private static void WriteParameterFramerate(Stream stream, Rational Framerate)
		{
			List<byte> data = new List<byte>();
			data.AddRange(ParaFramerate);
			data.AddRange(ToByteArray(Framerate.Numinator.ToString()));
			data.Add((byte)':');
			data.AddRange(ToByteArray(Framerate.Denominator.ToString()));
			stream.Write(data.ToArray(), 0, data.Count);
		}

		private static void WriteParameterColorspace(Stream stream, Y4MColorSpace Colorspace)
		{
			byte[] data = ToByteArray(" "+Colorspace.ToString());
			stream.Write(data, 0, data.Length);
		}
		#endregion

		int width;
		int height;
		Rational framerate;
		Y4MColorSpace colorspace;
		KrKbConstants KrKbStandard;
		bool headerWritten;

		public static long CalculateVideoSize(int Width, int Height, Rational Framerate, int NumberOfFrames, Y4MColorSpace Colorspace)
		{
			long HeaderSize;
			using (MemoryStream ms = new MemoryStream())
			{
				WriteHeader(ms, Width, Height, Framerate, Colorspace);
				HeaderSize = ms.Length;
				ms.Close();
			}
	
			long SizePerFrame = HeaderFrame.Length + 1;
			if (Colorspace == Y4MColorSpace.C420)
			{
				SizePerFrame += Width * Height * 3 / 2 + 1;
			}
			else if (Colorspace == Y4MColorSpace.C422)
			{
				SizePerFrame += Width * Height * 2;
			}
			else if (Colorspace == Y4MColorSpace.C444)
			{
				SizePerFrame += Width * Height * 3;
			}
			else
			{
				throw new ArgumentException("Invalid Colorspace", "Colorspace");
			}
			return HeaderSize + NumberOfFrames * SizePerFrame;
		}

		public static long CalculateVideoSize(int Width, int Height, Rational Framerate, TimeSpan Duration, Y4MColorSpace Colorspace)
		{
			int NumberOfFrames = (int)(Duration.TotalSeconds * Framerate.Numinator / Framerate.Denominator + 1);
			return CalculateVideoSize(Width, Height, Framerate, NumberOfFrames, Colorspace);
		}

		public Y4MEncoder(int Width, int Height, Rational Framerate, Y4MColorSpace Colorspace, KrKbConstants KrKbStandard)
		{
			this.width = Width;
			this.height = Height;
			this.framerate = Framerate;
			this.colorspace = Colorspace;
			this.KrKbStandard = KrKbStandard;
			this.headerWritten = false;
		}

		public bool AddImage(WriteableBitmap image, Stream stream)
		{
			if (image.PixelWidth != this.width || image.PixelHeight != this.height)
				return false;

			if (this.headerWritten == false)
			{
				WriteHeader(stream, this.width, this.height, this.framerate, this.colorspace);
				this.headerWritten = true;
			}

			stream.Write(HeaderFrame, 0, HeaderFrame.Length);
			stream.WriteByte((byte)0x0A);
			if (this.colorspace == Y4MColorSpace.C420)
				this.AddImage_C420(image, stream);
			else if (this.colorspace == Y4MColorSpace.C444)
				this.AddImage_C444(image, stream);

			return true;
		}

		private void AddImage_C420(WriteableBitmap image, Stream stream)
		{
			List<byte> Y = new List<byte>();
			List<byte> Cb = new List<byte>();
			List<byte> Cr = new List<byte>();

			for (int i = 0; i < image.Pixels.Length; i++)
			{
				Color c = new Color();
				byte[] YCbCr = ColorspaceConverter.Rgb_To_YCbCr(c.FromInt(image.Pixels[i]), this.KrKbStandard, 0);
				Y.Add(YCbCr[0]);

				int x = i % this.width;
				int y = i / this.width;

				if (x % 2 == 1 && y % 2 == 1)
				{
					Cb.Add(YCbCr[1]);
					Cr.Add(YCbCr[2]);
				}
			}
			stream.Write(Y.ToArray(), 0, Y.Count);
			stream.Write(Cb.ToArray(), 0, Cb.Count);
			stream.Write(Cr.ToArray(), 0, Cr.Count);
		}

		private void AddImage_C444(WriteableBitmap image, Stream stream)
		{
			List<byte> Y = new List<byte>();
			List<byte> Cb = new List<byte>();
			List<byte> Cr = new List<byte>();

			for (int i = 0; i < image.Pixels.Length; i++)
			{
				Color c = new Color();
				byte[] YCbCr = ColorspaceConverter.Rgb_To_YCbCr(c.FromInt(image.Pixels[i]), this.KrKbStandard, 0);
				Y.Add(YCbCr[0]);
				Cb.Add(YCbCr[1]);
				Cr.Add(YCbCr[2]);
			}
			stream.Write(Y.ToArray(), 0, Y.Count);
			stream.Write(Cb.ToArray(), 0, Cb.Count);
			stream.Write(Cr.ToArray(), 0, Cr.Count);
		}
	}
}
